import numpy as np
from scipy import signal
import time

def tm_ccoeff_normed_numpy(image, template):
    """
    使用NumPy实现cv2.TM_CCOEFF_NORMED模板匹配算法
    
    参数:
    - image: 原始图像 (灰度图)
    - template: 模板图像 (灰度图)
    
    返回:
    - 归一化互相关结果矩阵
    """
    # 获取图像和模板的尺寸
    h, w = template.shape
    H, W = image.shape
    
    # 1. 将image和template转换为float64，提高精度
    image = image.astype(np.float64)
    template = template.astype(np.float64)
    
    # 2. 计算模板的均值和移除均值
    template_mean = np.mean(template)
    template_centered = template - template_mean
    
    # 3. 使用快速卷积计算匹配结果
    # 使用FFT卷积，比直接卷积快很多，尤其对于大图像
    corr = signal.fftconvolve(image, template_centered[::-1, ::-1], mode='valid')
    
    # 4. 计算图像的局部均值 (使用卷积实现的滑动窗口)
    # 创建与模板相同大小的全1数组作为卷积核
    ones_template = np.ones_like(template, dtype=np.float64)
    
    # 使用FFT卷积快速计算图像每个滑动窗口的总和
    local_sum = signal.fftconvolve(image, ones_template[::-1, ::-1], mode='valid')
    
    # 计算每个滑动窗口的均值
    local_mean = local_sum / (h * w)
    
    # 5. 计算图像的局部标准差
    # 先计算每个像素的平方
    image_squared = image ** 2
    
    # 使用卷积计算滑动窗口内像素平方和
    local_sum_squared = signal.fftconvolve(image_squared, ones_template[::-1, ::-1], mode='valid')
    
    # 计算标准差: sqrt(E[X²] - (E[X])²)
    local_var = local_sum_squared / (h * w) - local_mean ** 2
    local_std = np.sqrt(np.maximum(local_var, 0))  # 使用maximum防止数值误差导致负数
    
    # 6. 计算模板的标准差
    template_std = np.std(template)
    
    # 7. 计算归一化互相关系数
    # 分子: (image - local_mean) 与 (template - template_mean) 的相关性
    # 分母: 两者的标准差乘积
    denominator = local_std * template_std * (h * w)
    
    # 处理除零问题
    is_zero_denominator = (denominator == 0)
    denominator[is_zero_denominator] = 1
    
    # 计算最终的归一化相关系数
    result = corr / denominator
    
    # 处理分母为0的位置
    result[is_zero_denominator] = 0
    
    return result

def tm_ccoeff_normed_numpy_optimized(image, template):
    """
    高度优化的cv2.TM_CCOEFF_NORMED NumPy实现
    
    使用内存优化和向量化操作来加速计算
    
    参数:
    - image: 原始图像 (灰度图)
    - template: 模板图像 (灰度图)
    
    返回:
    - 归一化互相关结果矩阵
    """
    # 获取图像和模板的尺寸
    h, w = template.shape
    H, W = image.shape
    
    # 结果矩阵的大小
    result_height = H - h + 1
    result_width = W - w + 1
    
    # 优化1: 提前转换数据类型，只转换一次
    image = image.astype(np.float32)
    template = template.astype(np.float32)
    
    # 优化2: 模板预处理 - 计算一次，重复使用
    template_mean = np.mean(template)
    template_centered = template - template_mean
    template_std = np.std(template)
    template_size = h * w
    
    # 优化3: 使用滑动窗口方法计算
    # 分配结果数组
    result = np.zeros((result_height, result_width), dtype=np.float32)
    
    # 优化4: 使用快速傅里叶变换进行快速卷积
    # 1. 计算图像和模板的加权相关性
    corr = signal.fftconvolve(image, template_centered[::-1, ::-1], mode='valid')
    
    # 2. 计算图像的局部均值
    local_sum = np.zeros((result_height, result_width), dtype=np.float32)
    box_filter = np.ones((h, w), dtype=np.float32)
    local_sum = signal.fftconvolve(image, box_filter[::-1, ::-1], mode='valid')
    local_mean = local_sum / template_size
    
    # 3. 计算图像的局部标准差
    image_squared = image ** 2
    local_sum_squared = signal.fftconvolve(image_squared, box_filter[::-1, ::-1], mode='valid')
    local_variance = (local_sum_squared / template_size) - (local_mean ** 2)
    # 确保方差非负 (避免数值精度问题)
    local_variance = np.maximum(local_variance, 0)
    local_std = np.sqrt(local_variance)
    
    # 4. 计算归一化相关系数
    # 处理除零问题
    denominator = local_std * template_std * template_size
    # 创建掩码标识除数为0的位置
    zero_mask = (denominator < 1e-10)
    # 避免除零
    denominator = np.where(zero_mask, 1, denominator)
    
    # 计算结果
    result = corr / denominator
    # 将除数为0的位置设为0
    result[zero_mask] = 0
    
    return result

# 更快的实现版本，使用滑动窗口和优化的FFT
def fast_tm_ccoeff_normed(image, template):
    """
    超快速实现的cv2.TM_CCOEFF_NORMED
    
    使用FFTW优化的实现，专注于大图像和模板的快速处理
    
    参数:
    - image: 原始图像 (灰度图, np.float32)
    - template: 模板图像 (灰度图, np.float32)
    
    返回:
    - 归一化互相关结果矩阵
    """
    # 强制转换为float32以提高速度
    if image.dtype != np.float32:
        image = image.astype(np.float32)
    if template.dtype != np.float32:
        template = template.astype(np.float32)
    
    # 获取尺寸
    h, w = template.shape
    template_size = h * w
    H, W = image.shape
    
    # 优化1: 模板前置处理
    template_mean = np.mean(template)
    template_centered = template - template_mean
    template_std = np.std(template)
    
    # 优化2: 使用预计算的FFT进行快速互相关
    # 将模板翻转180度
    template_flipped = template_centered[::-1, ::-1]
    
    # 使用快速傅里叶变换进行卷积
    # mode='valid'确保只返回完全重叠的部分
    numerator = signal.fftconvolve(image, template_flipped, mode='valid')
    
    # 优化3: 使用积分图像加速均值和标准差计算
    # 积分图像优化方案 - 在小模板大图像时特别有效
    if template_size > 1000:  # 只在模板较大时使用积分图像优化
        # 创建积分图像
        integral = np.cumsum(np.cumsum(image, axis=0), axis=1)
        integral_squared = np.cumsum(np.cumsum(image**2, axis=0), axis=1)
        
        # 初始化结果数组
        result_h, result_w = H-h+1, W-w+1
        local_mean = np.zeros((result_h, result_w), dtype=np.float32)
        local_std = np.zeros((result_h, result_w), dtype=np.float32)
        
        # 使用积分图像计算局部均值和标准差
        for i in range(result_h):
            for j in range(result_w):
                # 计算窗口的四个角
                x1, y1 = j, i
                x2, y2 = j+w, i+h
                
                # 计算窗口内像素和
                if x1 == 0 and y1 == 0:
                    window_sum = integral[y2-1, x2-1]
                elif x1 == 0:
                    window_sum = integral[y2-1, x2-1] - integral[y1-1, x2-1]
                elif y1 == 0:
                    window_sum = integral[y2-1, x2-1] - integral[y2-1, x1-1]
                else:
                    window_sum = integral[y2-1, x2-1] - integral[y1-1, x2-1] - integral[y2-1, x1-1] + integral[y1-1, x1-1]
                
                # 计算窗口内像素平方和
                if x1 == 0 and y1 == 0:
                    window_sum_squared = integral_squared[y2-1, x2-1]
                elif x1 == 0:
                    window_sum_squared = integral_squared[y2-1, x2-1] - integral_squared[y1-1, x2-1]
                elif y1 == 0:
                    window_sum_squared = integral_squared[y2-1, x2-1] - integral_squared[y2-1, x1-1]
                else:
                    window_sum_squared = integral_squared[y2-1, x2-1] - integral_squared[y1-1, x2-1] - integral_squared[y2-1, x1-1] + integral_squared[y1-1, x1-1]
                
                # 计算均值和标准差
                local_mean[i, j] = window_sum / template_size
                var = (window_sum_squared / template_size) - (local_mean[i, j] ** 2)
                local_std[i, j] = np.sqrt(max(0, var))
    else:
        # 对于小模板，使用FFT卷积可能更快
        box_filter = np.ones((h, w), dtype=np.float32)
        local_sum = signal.fftconvolve(image, box_filter[::-1, ::-1], mode='valid')
        local_mean = local_sum / template_size
        
        # 计算局部标准差
        local_sum_squared = signal.fftconvolve(image**2, box_filter[::-1, ::-1], mode='valid')
        local_var = (local_sum_squared / template_size) - (local_mean ** 2)
        local_std = np.sqrt(np.maximum(local_var, 0))
    
    # 计算分母
    denominator = local_std * template_std * template_size
    
    # 处理除零问题
    mask = denominator < 1e-8
    denominator[mask] = 1
    
    # 计算归一化相关系数
    result = numerator / denominator
    result[mask] = 0
    
    return result

# 速度测试函数
def benchmark_template_matching(image_size, template_size, implementation):
    """测试模板匹配算法的速度"""
    # 创建随机图像和模板
    image = np.random.randint(0, 256, (image_size, image_size)).astype(np.uint8)
    template = np.random.randint(0, 256, (template_size, template_size)).astype(np.uint8)
    
    # 预热
    implementation(image, template)
    
    # 测速
    start_time = time.time()
    result = implementation(image, template)
    end_time = time.time()
    
    return end_time - start_time, result.shape

# 使用示例
if __name__ == "__main__":
    # 创建测试图像和模板
    image_size = 1000
    template_sizes = [10, 50, 100, 200]
    
    print("速度测试不同实现的TM_CCOEFF_NORMED:")
    print("-" * 70)
    print(f"{'模板大小':<10} | {'基础实现(秒)':<15} | {'优化实现(秒)':<15} | {'超快实现(秒)':<15}")
    print("-" * 70)
    
    for template_size in template_sizes:
        image = np.random.randint(0, 256, (image_size, image_size)).astype(np.uint8)
        template = np.random.randint(0, 256, (template_size, template_size)).astype(np.uint8)
        
        # 基础实现
        start = time.time()
        tm_ccoeff_normed_numpy(image, template)
        basic_time = time.time() - start
        
        # 优化实现
        start = time.time()
        tm_ccoeff_normed_numpy_optimized(image, template)
        optimized_time = time.time() - start
        
        # 超快实现
        start = time.time()
        fast_tm_ccoeff_normed(image, template)
        fast_time = time.time() - start
        
        print(f"{template_size:<10} | {basic_time:<15.4f} | {optimized_time:<15.4f} | {fast_time:<15.4f}")
        
    print("-" * 70)
    print("注: 实际性能取决于图像和模板大小，以及硬件配置")